#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-07-24


import threading
import time
import cv2
import os
import json
import argparse
import numpy as np
from typing import Union
from queue import Queue
import platform
from spirems import tensor2sms, sms2tensor
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg, pcl2sms, BaseNode, get_extra_args
import base64
import copy
import matplotlib.pyplot as plt


class NuFoxgloveVisNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.imshow = self.get_param("imshow", True)
        self.sample_topic = self.get_param("sample_topic", "")
        self.resize = self.get_param("resize", [-1, -1])
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        input_topic = self.sample_topic if len(self.sample_topic) > 0 else '/' + job_name + '/sensor/sample'

        self._sample_reader = Subscriber(
            input_topic, 'spirecv_msgs::NuSample', self.sample_callback,
            ip=ip, port=port
        )
        self._image_writer = Publisher(
            '/' + job_name + '/detection_vis', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )

        self.sample_queue = Queue()
        self.queue_pool.append(self.sample_queue)
        self.foxglove_img_pubs = {}
        self.start()

    def release(self):
        BaseNode.release(self)
        self._sample_reader.kill()
        self._image_writer.kill()

    def sample_callback(self, msg):
        if self.sample_queue.empty():
            for sd in msg["sample_data"]:
                chan = sd["channel"]
                sd["data"] = sms2tensor(sd["data"])
            self.sample_queue.put(msg)

    def run(self):
        while self.is_running():
            sp_msg = self.sample_queue.get(block=True)
            if sp_msg is None:
                break

            # print(time.time() - sp_msg["t1"])

            for sd in sp_msg["sample_data"]:
                chan = sd["channel"]
                if chan in ['CAM_FRONT', 'CAM_FRONT_LEFT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK_RIGHT', 'CAM_BACK']:
                    if chan not in self.foxglove_img_pubs:
                        self.foxglove_img_pubs[chan] = Publisher('/to_foxglove/' + chan, 'sensor_msgs::CompressedImage')
                    img_msg = cvimg2sms(sd["data"], frame_id=chan)
                    self.foxglove_img_pubs[chan].publish(img_msg)
                if chan in ['LIDAR_TOP', 'LIDAR_FRONT', 'LIDAR_LEFT', 'LIDAR_RIGHT', 'LIDAR_FUSED']:
                    if chan not in self.foxglove_img_pubs:
                        self.foxglove_img_pubs[chan] = Publisher('/to_foxglove/' + chan, 'sensor_msgs::PointCloud')
                    lidar_msg = pcl2sms(sd["data"], sd["fields"], frame_id=chan)
                    self.foxglove_img_pubs[chan].publish(lidar_msg)
                if chan in ['RADAR_FRONT', 'RADAR_FRONT_RIGHT', 'RADAR_FRONT_LEFT', 'RADAR_BACK_LEFT', 'RADAR_BACK_RIGHT']:
                    if chan not in self.foxglove_img_pubs:
                        self.foxglove_img_pubs[chan] = Publisher('/to_foxglove/' + chan, 'sensor_msgs::PointCloud')
                    lidar_msg = pcl2sms(sd["data"].astype(np.float32), sd["fields"], frame_id=chan)
                    self.foxglove_img_pubs[chan].publish(lidar_msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = NuFoxgloveVisNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()

